/*
 * @lc app=leetcode.cn id=725 lang=cpp
 *
 * [725] 分隔链表
 */

// @lc code=start
/**
 * Definition for singly-linked list.
 * struct ListNode {
 *     int val;
 *     ListNode *next;
 *     ListNode() : val(0), next(nullptr) {}
 *     ListNode(int x) : val(x), next(nullptr) {}
 *     ListNode(int x, ListNode *next) : val(x), next(next) {}
 * };
 */

// #include <iostream>
// #include <vector>
// using namespace std;

// struct ListNode {
//     int val;
//     ListNode* next;
//     ListNode() : val(0), next(nullptr) {}
//     ListNode(int x) : val(x), next(nullptr) {}
//     ListNode(int x, ListNode* next) : val(x), next(next) {}
// };

class Solution {
   public:
    vector<ListNode*> splitListToParts(ListNode* head, int k) {
        int length = 0;                 // 保存初始链表的长度
        vector<ListNode*> result;       // 保存输出结果
        ListNode* curNode = head;
        ListNode* preNode = nullptr;

        // 遍历链表获取其长度
        while (curNode != nullptr) {
            curNode = curNode->next;
            length++;
        }

        int width = length / k;
        int remain = length % k;
        curNode = head;

        while (result.size() < k) {
            int tmpLength = (result.size() < remain) ? width + 1 : width;

            if (tmpLength == 0) {
                result.emplace_back(nullptr);
            } else {
                result.emplace_back(curNode);
                while (tmpLength-- > 0) {
                    preNode = curNode;
                    curNode = curNode->next;
                }
                preNode->next = nullptr;
            }
        }

        return result;
    }
};
// @lc code=end
